from typing import Optional

from manimlib import *

class CodeLines(VGroup):
    def __init__(self, text: str, t2c={}) -> None:
        super().__init__()

        # 将文本分为各行，并创建 Text 对象
        for text in text.splitlines():
            self.add(Text(text, base_color=GREY_B, t2c=t2c, font='Source Code Pro Medium').scale(0.4))
        
        self.arrange(DOWN, aligned_edge=LEFT, buff=SMALL_BUFF)

    def replaceAnimate(self, fn: Callable, index: int, text: str, t2c={}, animate=ReplacementTransform) -> None:
        new = Text(text, base_color=GREY_B, t2c=t2c).scale(0.4)
        old = self[index]
        new.move_to(old.get_left(), LEFT)
        fn(animate(old, new))
        self.replace_submobject(index, new)

    def removeAnimate(self, start: int, end: Optional[int]=None, animate=FadeOut) -> AnimationGroup:
        '''
        从代码中删除某（几）行的动画
        包括把删除部分之后的代码向上移动的动作
        '''
        if not end:
            end = start + 1
        sub = self[start:end]
        self.remove(*self[start:end])
        ul = self.get_corner(UL)
        return AnimationGroup(animate(sub), self.animate.arrange(DOWN, aligned_edge=LEFT, buff=SMALL_BUFF).move_to(ul, UL))

    def set_code_color(self, start: int, end: int, color: ManimColor, **kwargs) -> None:
        '''
        按照原文本区间设置文字颜色
        '''
        for txt in self.submobjects:    # 遍历每一行的文本
            length = len(txt.text) + 1      # 考虑一个换行符

            # 如果起止区间包含该行文本，则将对应区段设置颜色
            if start < length:
                v_start, v_end = start, end
                idx = 0
                while idx < start:
                    if txt.text[idx] == ' ':
                        v_start -= 1
                        v_end -= 1
                    idx += 1
                while idx < end:
                    if txt.text[idx] == ' ':
                        v_end -= 1
                    idx += 1
                
                txt[v_start:v_end].set_color(color, **kwargs)
                start = 0
            else:
                start -= length
            end -= length

            # end <= 0，说明区段已经结束
            if end <= 0:
                break

class CodeBackground(BackgroundRectangle):
    CONFIG = {
        "fill_color": BLACK,
        "fill_opacity": 1,
        "stroke_width": 1,
        "stroke_opacity": 1,
        "stroke_color": GREY_A,
        "buff": 0.2
    }
class CodeHeader(Rectangle):
    CONFIG = {
        "fill_opacity": 1,
        "stroke_width": 1,
        "stroke_opacity": 1,
        "stroke_color": GREY_A,
    }

class CodePlane(VGroup):
    def __init__(self, texts, t2c) -> None:
        self.lines = CodeLines(texts, t2c)
        self.background = CodeBackground(self.lines)
        super().__init__(self.background, self.lines)

class CodeView(VGroup):
    def __init__(self, title, texts, t2c, title_color=WHITE, header_color=GREY_D) -> None:
        self.lines = CodeLines(texts, t2c)
        self.background = CodeBackground(self.lines)
        txt_title = Text(title, color = title_color).scale(0.35)
        self.header = CodeHeader(self.background.get_width(), txt_title.get_height() + 0.2, fill_color=header_color)
        self.header.move_to(self.background.get_top() + [0, self.header.get_height() / 2, 0])
        txt_title.move_to(self.header)
        self.others = VGroup(self.header, self.background, txt_title)
        super().__init__(self.others, self.lines)